"""
1.	机器人是模仿人类和动物行为的机器，其内部有台计算机，通过读取各个传感器的信息，做出判断，并且调用电机实现相关的动作，完成指令。
给定“手势识别”数据集，有三个指令：左转、右转、停止。利用keras深度学习平台预训练模型mobilenet，搭建后端网络，进行模型训练和测试，
按下面的要求，完成相应代码（56分）
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks
import numpy as np
import os
import sys
import cv2 as cv
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
import pandas as pd

np.random.seed(1)
tf.random.set_seed(1)

VER = 'v1.0'
SIZE = 224
ALPHA = 1e-3
BATCH_SIZE = 16
N_EPOCHS = 8
BASE_DIR, FILE_NAME = os.path.split(__file__)
dir = '../../../../large_data/CV4/_many_files/Gesture_Recognition'
IMG_DIR = os.path.join(BASE_DIR, dir)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

# ①	导入“手势识别”数据集
print('Loading images ...')
x, y, pathes, idx2label, label2idx = [], [], [], {}, {}
yi = 0
for sub_dir_name in os.listdir(IMG_DIR):
    idx2label[yi] = sub_dir_name
    label2idx[sub_dir_name] = yi
    sub_dir_path = os.path.join(IMG_DIR, sub_dir_name)
    for file_name in os.listdir(sub_dir_path):
        file_path = os.path.join(sub_dir_path, file_name)
        img = cv.imread(file_path, cv.IMREAD_COLOR)
        img = cv.resize(img, (SIZE, SIZE), interpolation=cv.INTER_CUBIC)
        x.append(img)
        y.append(yi)
        pathes.append(file_path)
    yi += 1
x = np.float32(x) / 255.
y = np.int64(y)
pathes = np.array(pathes)
n_cls = len(np.unique(y))
print('x', x.shape)
print('y', y.shape)
print('n_cls', n_cls)
print('Loaded')

# ②	按适当比例划分训练集、验证集、测试集
x_train, x_val_test, y_train, y_val_test, path_train, path_val_test = train_test_split(x, y, pathes, train_size=0.8, random_state=1, shuffle=True)
x_val, x_test, y_val, y_test, path_val, path_test = train_test_split(x_val_test, y_val_test, path_val_test, train_size=0.5, random_state=1, shuffle=True)
print('x_train', x_train.shape)
print('x_val', x_val.shape)
print('x_test', x_test.shape)
print('y_train', y_train.shape)
print('y_val', y_val.shape)
print('y_test', y_test.shape)
y_test_pd = pd.Series(y_test)
print(y_test_pd.value_counts())
print('path_train', path_train.shape)
print('path_val', path_val.shape)
print('path_test', path_test.shape)
dl_train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE)
dl_val = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE)
dl_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE)

# ③	创建模型：包括预训练模型mobilenet和后端网络，实现3种手势分类任务
base_model = keras.applications.mobilenet.MobileNet(
    input_shape=(SIZE, SIZE, 3),
    include_top=False,
    weights='imagenet',
    pooling='avg',
    classes=1000,
)
base_model.trainable = False
customer_model = layers.Dense(n_cls, activation=activations.softmax)
model = keras.Sequential([
    base_model,
    customer_model,
])
model.summary()

# ④	进行模型编译和训练，打印输出训练集、验证集、测试集准确率
model.compile(
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    loss=losses.sparse_categorical_crossentropy,
    metrics=[metrics.sparse_categorical_accuracy]
)
history = model.fit(
    dl_train,
    validation_data=dl_val,
    epochs=N_EPOCHS,
)
his = history.history
loss = his['loss']
val_loss = his['val_loss']
sparse_categorical_accuracy = his['sparse_categorical_accuracy']
val_sparse_categorical_accuracy = his['val_sparse_categorical_accuracy']

result = model.evaluate(dl_test)
test_loss = result[0]
test_acc = result[1]

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Loss')
plt.plot(loss, label='train')
plt.plot(val_loss, label='val')
plt.legend()
plt.grid()
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Accuracy')
plt.plot(sparse_categorical_accuracy, label='train')
plt.plot(val_sparse_categorical_accuracy, label='val')
plt.legend()
plt.grid()
plt.show()

# ⑤	模型训练完成后，用饼图显示预测正确和不正确所占比率
# ⑥	打印输出混淆矩阵，并且热力图显示混淆矩阵
pred = model.predict(dl_test)
pred = pred.argmax(axis=1)

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Accuracy')
plt.pie([test_acc, 1-test_acc], explode=[0.1, 0], labels=['Right', 'Wrong'], autopct='%0.2f%%')
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Confusion matrix')
mat = confusion_matrix(y_test, pred)
sns.heatmap(mat, annot=True)
plt.show()

# ⑦	从测试集中随机选取9张图片，利用训练模型，图示预测种类名称、并与真实种类名称比较，预测正确显示“黑色”，预测错误显示“红色”。
spr = 3
spc = 3
spn = 0
plt.figure(figsize=[8, 8])
for i in range(spr * spc):
    spn += 1
    plt.subplot(spr, spc, spn)
    predi = pred[i]
    yi = y_test[i]
    title = f'{idx2label[yi]}=>{idx2label[predi]}'
    plt.title(title, color='black' if predi == yi else 'red')
    img = cv.imread(path_test[i], cv.IMREAD_COLOR)[:, :, ::-1]
    plt.imshow(img)
plt.show()
